package com.lixw.langchain.controller;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.MetadataFilterBuilder;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

/**
 * @ClassName: EmbeddinglController
 * @Description:
 * @Author: xuweiLi
 * @Create: 2025/8/25 18:52
 **/
@Slf4j
@RestController
public class EmbeddingController {
    @Resource
    private EmbeddingModel embeddingModel;

    @Resource
    private QdrantClient qdrantClient;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    @GetMapping("/embedding/embed")
    public String embed() {
        String prompt = """
                《观沧海》
                
                东临碣石，以观沧海。
                
                水何澹澹，山岛竦峙。
                
                树木丛生，百草丰茂。
                
                秋风萧瑟，洪波涌起。
                
                日月之行，若出其中；
                
                星汉灿烂，若出其里。
                
                幸甚至哉，歌以咏志。
                """;
        //向量化数据
        Response<Embedding> embeddingResponse = embeddingModel.embed(prompt);
        log.info("向量化后的数据：{}", embeddingResponse); //1024长度的浮点数数组
        return embeddingResponse.content().toString();
    }

    @GetMapping("embedding/createCollection")
    public void createCollection() {
        Collections.VectorParams vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)
                .build();
        qdrantClient.createCollectionAsync("test-qdrant", vectorParams);
    }

    @GetMapping("embedding/add")
    public String add() {
        String prompt = """
                《观沧海》
                
                东临碣石，以观沧海。
                
                水何澹澹，山岛竦峙。
                
                树木丛生，百草丰茂。
                
                秋风萧瑟，洪波涌起。
                
                日月之行，若出其中；
                
                星汉灿烂，若出其里。
                
                幸甚至哉，歌以咏志。
                """;
        TextSegment segment = TextSegment.from(prompt);
        segment.metadata().put("author","曹操");

        Embedding embedding = embeddingModel.embed(segment).content();

        String addResult = embeddingStore.add(embedding, segment);//将向量化后的数据和原文保存起来

        log.info("<新增向量返回结果>: {}", addResult);

        return addResult;
    }

    @GetMapping("/embedding/query1")
    public void query1() {
        Embedding queryEmbedding = embeddingModel.embed("观沧海讲的是什么").content();

        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);

        log.info("<结果>: {}", searchResult);
        log.info("<查询向量数据库得到的结果>: {}", searchResult.matches().get(0).embedded().text());
    }

    @GetMapping("/embedding/query2")
    public void query2() {
        Embedding queryEmbedding = embeddingModel.embed("观沧海").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .filter(MetadataFilterBuilder.metadataKey("author").isEqualTo("曹操"))
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        if(searchResult.matches().isEmpty()){
            log.error("找不到查询作者这首诗");
            return;
        }
        log.info("<结果>: {}", searchResult);
        log.info("<查询向量数据库得到的结果>: {}", searchResult.matches().get(0).embedded().text());
    }

}